library(cjoint)
library(ggplot2)
library(data.table)
library(dplyr)
library(stringi)
library(RColorBrewer)
theme_set(theme_bw())

load("data_all.rda")

use.data <- data_all[!is.na(data_all$Chosen_party), ] %>%
  rename(psid = ID)

results <- amce(Chosen_party ~ Economy + Jobs + Immigration + EU + Education,
                data = use.data, design = ukdesign, cluster=F)

## Main model (Figure 2)
source("01_sub1_process_results.R")
ggsave("figure2.pdf", plot = g, width= 7, height = 5)

## Heterogeneity Model (Figure 3) 
source("01_sub2_process_het_model.R")
ggsave("figure3.pdf", plot = g, width = 8, height = 8.8) 


# Predictions (Figure 4)

use.data.3 <- use.data.2
for(i in 3:7) use.data.3[, i] <- use.data.3[, i] %>% as.numeric %>% factor
partypolicies <- gdata::read.xls("Policy_Table.xlsx")
for(i in 2:6) partypolicies[, i] <- partypolicies[, i] %>% factor %>% as.numeric %>% factor

output_plot_data_by_party <- list()
results <- lm(Chosen_party ~ Economy + Jobs + Immigration + EU + Education,
              data = use.data.3)
party <- "All"
pred <- predict(results, newdata = partypolicies, 
                interval = "confidence") %>% data.table
pred[, party_out := partypolicies$Pty]
pred[, party_support := party]
output_plot_data_by_party[[party]] <- pred

for(j in 1:5){
  # i <- c(2,3,4,6,8,9)[j]
  # party <- c("Cons", "Lab", "LibDem", "SNP", "UKIP", "Greens")[j]
  i <- c(2,3,4,8,9)[j]
  party <- c("Cons", "Lab", "LibDem", "UKIP", "Greens")[j]
  
  results <- lm(Chosen_party ~ Economy + Jobs + Immigration + EU + Education,
                  data = use.data.3[ which(use.data.3$voteint.2 == i), ])
  pred <- predict(results, newdata = partypolicies, 
                  interval = "confidence") %>% data.table
  pred[, party_out := partypolicies$Pty]
  pred[, party_support := party]
  output_plot_data_by_party[[party]] <- pred
}
output_plot_data_by_party <- rbindlist(output_plot_data_by_party)
names(output_plot_data_by_party) <- 
  c("point_pred", "ll", "ul", "Party_Manifesto", "voteint")
output_plot_data_by_party[, colour := ifelse(Party_Manifesto == voteint, "red", "black")]
output_plot_data_by_party$Party_Manifesto <- 
  output_plot_data_by_party$Party_Manifesto %>% 
  as.character %>% factor(levels = rev(c("Cons", "Lab", "LibDem", "UKIP", "Greens")))
output_plot_data_by_party$voteint <- 
  output_plot_data_by_party$voteint %>% 
  factor(levels = c("All", "Cons", "Lab", "LibDem", "UKIP", "Greens"))
ggplot(output_plot_data_by_party, aes(x = Party_Manifesto, y = point_pred,
                                      colour = colour)) + 
  geom_pointrange( aes(ymax = ul, ymin = ll), size = 0.2) + 
  scale_color_manual(values = c("grey50", "black")) +
  coord_flip() + facet_wrap(~ voteint) +
  xlab("Manifesto") +
  ylab("Predicted Preference") + ylim(c(-0.03,1.03)) + 
  theme(legend.position = "none")
ggsave("figure4.pdf", width = 7, height = 4)

